import logging
import os

import gradio as gr
import httpx
from langchain.agents import AgentExecutor, create_openai_tools_agent
from langchain.callbacks import FileCallbackHandler
from langchain.tools import Tool
from langchain_community.llms import Ollama
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_openai import ChatOpenAI
from tenacity import retry, stop_after_attempt, wait_exponential

# 配置参数
os.environ["LANGUAGE"] = "en_US"
DEEPSEEK_BASE_URL = "https://api.modelarts-maas.com/v1"
MODEL_NAME = "DeepSeek-V3"
OLLAMA_MODEL = "llama3"  # 备用模型
REQUEST_TIMEOUT = 10  # 请求超时时间（秒）

# 配置日志回调
logging.basicConfig(filename='agent.log', level=logging.INFO)
file_callback = FileCallbackHandler('agent.log')

# === 1. 优化 HTTP 客户端 ===
client = httpx.Client(timeout=REQUEST_TIMEOUT)


# === 2. 带重试机制的 DeepSeek 初始化 ===
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10))
def init_deepseek_llm() -> ChatOpenAI:
    """初始化 DeepSeek LLM，自动重试失败请求"""
    llm = ChatOpenAI(
        model=MODEL_NAME,
        temperature=0.7,
        base_url=DEEPSEEK_BASE_URL,
        http_client=client,  # 复用连接
        max_retries=2
    )
    # 测试连接
    llm.invoke("Ping")  # 轻量测试请求
    return llm


# === 3. 初始化工具 ===

def calculate(math_expression: str) -> str:
    """安全计算数学表达式"""
    try:
        # 限制可用的数学操作
        allowed_chars = set("0123456789+-*/.() ")
        if not all(c in allowed_chars for c in math_expression):
            raise ValueError("仅支持基础数学运算")
        return str(eval(math_expression))
    except Exception as e:
        return f"计算错误: {str(e)}"


tools = [
    Tool(
        name="Calculator",
        func=calculate,
        description="计算数学表达式，如 '(3 + 5) * 2'。仅支持基础运算。"
    )
]


# === 4. 带缓存的 Agent 创建 ===
def create_agent() -> AgentExecutor:
    """创建带缓存的 Agent 实例"""
    # 优先尝试 DeepSeek
    try:
        llm = init_deepseek_llm()
        print("DeepSeek 服务已连接")
    except Exception as e:
        print(f"DeepSeek 不可用: {str(e)}, 使用备用模型 {OLLAMA_MODEL}")
        llm = Ollama(model=OLLAMA_MODEL)

    prompt = ChatPromptTemplate.from_messages([
        ("system", "你是有用的AI助手，可以使用工具。如果用户问题涉及敏感信息，请拒绝回答。"),
        ("user", "{input}"),
        MessagesPlaceholder("agent_scratchpad"),
    ])

    agent = create_openai_tools_agent(llm, tools, prompt)
    return AgentExecutor(
        agent=agent,
        tools=tools,
        verbose=True,
        handle_parsing_errors=True,  # 自动处理解析错误
        max_iterations=5  # 限制最大迭代次数
    )


# 全局 Agent 实例
agent_executor = create_agent()


# === 5. 优化请求处理 ===

def run_agent(query: str, chat_history: list) -> tuple[str, list]:
    if not query.strip():
        return "", chat_history + [(query, "请输入有效问题")]

    try:
        response = agent_executor.invoke(
            {"input": query},
            config={"callbacks": [file_callback]}  # 使用修正后的回调
        )
        return "", chat_history + [(query, response["output"])]
    except Exception as e:
        error_msg = f"服务暂时不可用: {str(e)}" if "rate limit" not in str(e).lower() else "请求过于频繁"
        return "", chat_history + [(query, error_msg)]


# === 6. 优化的 Gradio 界面 ===
with gr.Blocks(title="AI 助手 (DeepSeek 优化版)") as demo:
    gr.Markdown("""
    ## 🚀 DeepSeek 智能助手
    - 支持数学计算
    - 自动故障转移至本地模型
    """)

    with gr.Row():
        chatbot = gr.Chatbot(height=400, label="对话历史")

    with gr.Row():
        msg = gr.Textbox(
            label="输入问题",
            placeholder="输入后按回车...",
            max_lines=3
        )
        clear_btn = gr.Button("清空", variant="secondary")

    # 交互逻辑
    msg.submit(
        run_agent,
        inputs=[msg, chatbot],
        outputs=[msg, chatbot],
        queue=True  # 启用队列防止并发问题
    )
    clear_btn.click(
        lambda: None,
        None,
        chatbot,
        queue=False
    )

# === 7. 启动优化 ===
if __name__ == "__main__":
    try:
        demo.launch(
            server_name="0.0.0.0",
            server_port=7860,
            favicon_path="./favicon.ico",  # 可选: 自定义图标
            share=False,
            show_error=True  # 显示详细错误
        )
    except KeyboardInterrupt:
        client.close()  # 确保关闭 HTTP 连接
        print("服务已安全停止")
